﻿using System;
using System.IO;
using System.Linq;
using System.Text;
using MLSharp.Classification;

namespace MLSharp.Evaluation
{
	/// <summary>
	/// Creates a grid of confusion matrices at different
	/// thresholds, which can then be used to create an ROC 
	/// curve.
	/// </summary>
	public class ConfusionGrid
	{
		#region Private Fields

		/// <summary>
		/// The confusion matrices for each threshold arranged from 
		/// 0 to 1.
		/// </summary>
		private readonly ConfusionMatrix[] mMatrices;

		#endregion

		#region Public Constructors

		/// <summary>
		/// Creates a confusion grid using the classification results of 
		/// the specified dataset.
		/// </summary>
		/// <param name="dataSet"></param>
		/// <param name="results"></param>
		public ConfusionGrid(IDataSet dataSet, ClassificationResult[] results)
		{
			if (dataSet.TargetAttributeIndex < 0)
			{
				throw new InvalidOperationException("Dataset's target attribute is not set correctly.");
			}

			if (dataSet.Attributes[dataSet.TargetAttributeIndex].Type != AttributeType.Set ||
				dataSet.Attributes[dataSet.TargetAttributeIndex].PossibleValues.Length != 2)
			{
				throw new InvalidOperationException("Only binary target attributes are supported.");
			}

			//Grab a copy of the instances list so we can reorder it.
			Instance[] instances = dataSet.Instances.ToArray();

			string negativeValue = dataSet.Attributes[dataSet.TargetAttributeIndex].PossibleValues[0];
			string positiveValue = dataSet.Attributes[dataSet.TargetAttributeIndex].PossibleValues[1];

			//Create a normalized array of the confidences.  Negative predictions
			//are given a normalized confidence in the [0,0.5] range and positive
			//predictions are given a normalized confidence in the [0.5,1] range.
			double[] normalizedConfidence =
				results.Select(result => result.PredictedClass == positiveValue ? result.Confidence : 1.0 - result.Confidence).ToArray();

			//Loop through and create a confusion matrix for each threshold in [0,1] at
			//0.1 increments.
			mMatrices = new ConfusionMatrix[11];

			for (int i=0; i < mMatrices.Length; i++)
			{
				double threshold = i*0.1;

				mMatrices[i] = new ConfusionMatrix(negativeValue, positiveValue);

				//Loop through the data and accumulate counts.
				for (int j=0; j < instances.Length; j++)
				{
					double confidence = normalizedConfidence[j];

					//If the confidence is greater than or equal to 
					//the threshold, count it as predicting true.
					if (confidence > threshold)
					{
						if (results[j].ActualClass == positiveValue)
						{
							mMatrices[i].TruePositives++;
						}
						else if (results[j].ActualClass == negativeValue)
						{
							mMatrices[i].FalsePositives++;
						}
						else
						{
							throw new InvalidOperationException("Unexpected value: " + instances[j].ClassValue);
						}
					}
					//Otherwise, count it as a negative prediction.
					else
					{
						if (results[j].ActualClass == negativeValue)
						{
							mMatrices[i].TrueNegatives++;
						}
						else if (results[j].ActualClass == positiveValue)
						{
							mMatrices[i].FalseNegatives++;
						}
						else
						{
							throw new InvalidOperationException("Unexpected value: " + instances[j].ClassValue);
						}
					}
				}
			}
		}

		#endregion

		#region Public Methods

		/// <summary>
		/// Gets the confusion matrix at the specified index.
		/// </summary>
		/// <param name="thresholdIndex"></param>
		public ConfusionMatrix GetStats(int thresholdIndex)
		{
			return mMatrices[thresholdIndex];
		}

		/// <summary>
		/// Outputs a list of grids, one per threshold.
		/// </summary>
		/// <returns></returns>
		public override string ToString()
		{
			StringBuilder sb = new StringBuilder();

			for (int i=0; i < mMatrices.Length; i++)
			{
			    sb.AppendFormat("Threshold {0}:", i*0.1).AppendLine();
			    sb.Append(mMatrices[i]);
			}

			return sb.ToString();
		}

		/// <summary>
		/// Writes the grid to an HTML file.  This allows the grid to be 
		/// easily used in Excel for more advanced analysis.
		/// </summary>
		/// <param name="filename"></param>
		public void SaveHtmlFile(string filename)
		{
			using (StreamWriter writer = File.CreateText(filename))
			{
				//Write usual HTML header stuff.
				writer.WriteLine("<html>\r\n<head></head>\r\n<body>");

				//Table open!
				writer.WriteLine("<table>");

                //Write one row for the threshold headers
				writer.WriteLine("<tr>");
				writer.WriteLine("<td>Thresholds:</td>");
				
				for (int i=0; i < mMatrices.Length; i++)
				{
					writer.WriteLine("<td colspan=\"2\">{0}</td>", i*0.1);
				}

				writer.WriteLine("</tr>");

				//Write a second row for the actual label headers
				writer.WriteLine("<tr>");

				writer.WriteLine("<td>&nbsp;</td>");
				
				for (int i=0; i < mMatrices.Length; i++)
				{
					writer.WriteLine("<td>{0}</td>", mMatrices[0].PositiveLabel);
					writer.WriteLine("<td>{0}</td>", mMatrices[0].NegativeLabel);
				}

				writer.WriteLine("</tr>");

				//Write two rows for the confusion matrices.
				//True Positives and False Positives first!
				writer.WriteLine("<tr>");

				writer.WriteLine("<td>predicted-{0}</td>", mMatrices[0].PositiveLabel);

				for (int i=0; i < mMatrices.Length; i++)
				{
					writer.WriteLine("<td>{0}</td>", mMatrices[i].TruePositives);
					writer.WriteLine("<td>{0}</td>", mMatrices[i].FalsePositives);
				}

				writer.WriteLine("</tr>");

				//False Negatives and True Negatives next!

				writer.WriteLine("<tr>");

				writer.WriteLine("<td>predicted-{0}</td>", mMatrices[0].NegativeLabel);

				for (int i = 0; i < mMatrices.Length; i++)
				{
					writer.WriteLine("<td>{0}</td>", mMatrices[i].FalseNegatives);
					writer.WriteLine("<td>{0}</td>", mMatrices[i].TrueNegatives);
				}

				writer.WriteLine("</tr>");

				//Close out the table
				writer.WriteLine("</table>");
				writer.WriteLine("</body>");
			}
		}

		#endregion
	}
}
